package com.study.langchain4j.config.memory;

import com.study.langchain4j.util.RedisOpcTemplate;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.ObjectUtils;

import java.util.List;

/**
 * @author xl-9527
 * @since 2025/10/30
 **/
@Configuration(proxyBeanMethods = false)
public class MemoryConfig {

    private static final Logger log = LogManager.getLogger(MemoryConfig.class);

    private static final String MEMORY_ID = "ai:memory_id:";

    @Bean
    ChatMemoryProvider messageWindowMemoryProvider(final ChatMemoryStore chatMemoryStore) {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(100)
                .chatMemoryStore(chatMemoryStore)
                .build();
    }

    @Bean
    ChatMemoryProvider tokenWindowChatMemory() {
        return memoryId -> TokenWindowChatMemory.withMaxTokens(
                100,
                new OpenAiTokenCountEstimator("qwen3-max")
        );
    }

    @Bean
    ChatMemoryStore redisChatMemoryStore(final RedisOpcTemplate redisOpcTemplate) {
        return new ChatMemoryStore() {

            @Override
            public List<ChatMessage> getMessages(final Object memoryId) {
                final String json = redisOpcTemplate.get(MEMORY_ID + memoryId.toString());
                if (ObjectUtils.isEmpty(json)) {
                    return List.of();
                }
                return ChatMessageDeserializer.messagesFromJson(json);
            }

            @Override
            public void updateMessages(final Object memoryId, final List<ChatMessage> messages) {
                redisOpcTemplate.set(MEMORY_ID + memoryId.toString(), ChatMessageSerializer.messagesToJson(messages));
            }

            @Override
            public void deleteMessages(final Object memoryId) {
                if (log.isDebugEnabled()) {
                    log.debug("delete memoryId: {}", memoryId);
                }
                redisOpcTemplate.delete(MEMORY_ID + memoryId.toString());
            }
        };
    }
}
